import jieba
import wordcloud
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from collections import Counter
import os


def read_stopwords(file_path="D:\\嘟嘟\\历史数据\\python学习\\stopwords.txt"):
    """读取停用词表"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return set([line.strip() for line in f])


def read_text(file_path):
    """读取文本文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return f.read()


def process_text(text, stopwords):
    """处理文本，进行分词并去除停用词"""
    words = jieba.cut(text)
    return ' '.join([word for word in words if word not in stopwords and len(word) > 1])


def generate_wordcloud(text, font_path, mask=None, width=800, height=400, max_words=100, background_color='blue'):
    """生成词云图"""
    wc = wordcloud.WordCloud(
        font_path=font_path,
        width=width,
        height=height,
        max_words=max_words,
        background_color=background_color,
        mask=mask
    )
    wc.generate(text)
    return wc


def create_circle_mask(size=800):
    """创建圆形蒙版"""
    mask = np.zeros((size, size), dtype=np.uint8)
    center = size // 2
    radius = size // 2 - 10
    y, x = np.ogrid[:size, :size]
    dist_from_center = np.sqrt((x - center) ** 2 + (y - center) ** 2)
    mask[dist_from_center <= radius] = 255
    return mask


def main():
    # 读取文本和停用词
    text = read_text("D:\\嘟嘟\\历史数据\\python学习\\实践项目1 jieba分词\\绘制词云图用文档\\唐诗\\唐诗三百首.txt")
    stopwords = read_stopwords()
    processed_text = process_text(text, stopwords)

    # 字体路径
    fonts = {
        '微软雅黑': "D:\\嘟嘟\\历史数据\\python学习\\msyh.ttc",
        '微软雅黑粗体':"D:\\嘟嘟\\历史数据\\python学习\\msyhbd.ttc",
        '微软雅黑细体': "D:\\嘟嘟\\历史数据\\词云作业计科233黄煜2023081313\\msyhl.ttc"
    }

    # 生成不同尺寸的词云图
    sizes = [(800, 400), (1000, 500), (1200, 600)]
    for width, height in sizes:
        wc = generate_wordcloud(processed_text, fonts['微软雅黑'], width=width, height=height)
        plt.figure(figsize=(10, 5))
        plt.imshow(wc, interpolation='bilinear')
        plt.axis('off')
        plt.savefig(f'wordcloud_{width}x{height}.png')
        plt.close()

    # 生成不同关键词数量的词云图
    word_counts = [30, 50, 100]
    for count in word_counts:
        wc = generate_wordcloud(processed_text, fonts['微软雅黑'], max_words=count)
        plt.figure(figsize=(10, 5))
        plt.imshow(wc, interpolation='bilinear')
        plt.axis('off')
        plt.savefig(f'wordcloud_{count}words.png')
        plt.close()

    # 使用不同字体生成词云图
    for font_name, font_path in fonts.items():
        wc = generate_wordcloud(processed_text, font_path)
        plt.figure(figsize=(10, 5))
        plt.imshow(wc, interpolation='bilinear')
        plt.axis('off')
        plt.savefig(f'wordcloud_{font_name}.png')
        plt.close()

    # 使用不同形状生成词云图
    # 1. 五角星形状
    star_mask = np.array(Image.open('star_mask.png'))
    wc = generate_wordcloud(processed_text, fonts['微软雅黑'], mask=star_mask)
    plt.figure(figsize=(10, 10))
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')
    plt.savefig('wordcloud_star.png')
    plt.close()

    # 2. 圆形形状
    circle_mask = create_circle_mask()
    wc = generate_wordcloud(processed_text, fonts['微软雅黑'], mask=circle_mask)
    plt.figure(figsize=(10, 10))
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')
    plt.savefig('wordcloud_circle.png')
    plt.close()

    # 3. 镰刀形状 (如果存在镰刀形状的蒙版图片)
    try:
        sickle_mask = np.array(Image.open('sickle_mask.png'))
        wc = generate_wordcloud(processed_text, fonts['微软雅黑'], mask=sickle_mask)
        plt.figure(figsize=(10, 10))
        plt.imshow(wc, interpolation='bilinear')
        plt.axis('off')
        plt.savefig('wordcloud_sickle.png')
        plt.close()
    except FileNotFoundError:
        print("镰刀形状蒙版图片不存在，跳过生成镰刀形状词云图")


if __name__ == '__main__':
    main()